import cv2
import numpy as np

class MySimlarity:
    def __init__(self, thres):
        self.thres = thres

    def __call__(self, im, sp, tp):
        sx = sp[0]
        sy = sp[1]

        tx = tp[0]
        ty = tp[1]

        diff = abs(im[sy, sx] - im[ty, tx])

        return diff < self.thres


class RegionGrowing:
    def __init__(self, similarity_func):
        self.similarity_func = similarity_func
        self.near_point = [[1, 0], [1, 1], [0, 1], [0, -1], [-1, 0], [-1, -1], [-1, 1], [1, -1]]

    def __call__(self, src, x, y, map=None, lab=1):
        h, w = src.shape

        if map is None:
            map = np.zeros_like(src).tolist()

        stack = []
        area = []

        stack.append((x, y))
        area.append((x, y))
        map[x][y] = lab

        while len(stack):
            ox, oy = stack.pop(0)

            for d in self.near_point:
                nx = ox + d[0]
                ny = oy + d[1]

                if nx >= 0 and nx <= w - 1 and ny >= 0 and ny <= h - 1:
                    if map[nx][ny] == 0:
                        if self.similarity_func(src, (x, y), (nx, ny)):
                            stack.append((nx, ny))
                            area.append((nx, ny))
                            map[nx][ny] = lab

        return map, area


class GrowSegment:
    def __init__(self, regiongrow_algorithm):
        self.regiongrow_algorithm = regiongrow_algorithm

    def __call__(self, src, thres):
        
        h, w = src.shape

        map = np.zeros_like(src).tolist()

        lab = 1

        sections = []

        for y in range(0, h):
            for x in range(0, w):
                if map[y][x] == 0:
                    map, area = self.regiongrow_algorithm(src, x, y, map, lab)
                    lab += 1

                    if len(area) > thres:
                        sections.append(area)
        
        map = np.array(map).astype('uint8')

        segment = np.zeros_like(src).tolist()

        lab = 1
        for area in sections:
            for p in area:
                segment[p[1]][p[0]] = lab
            lab += 1

        segment = np.array(segment).astype('uint8')

        return map, segment
    

if __name__ == '__main__':
    SRC = cv2.imread('lena.jpg', cv2.IMREAD_GRAYSCALE)

    def show(winname, img):
        cv2.imshow(winname, img.astype('uint8'))
        cv2.resizeWindow(winname, 512, 512)
        cv2.waitKey(1)

    show('ImageGrowSegment', (GrowSegment(RegionGrowing(MySimlarity(100)))(SRC, 50)[1] * 4))
    cv2.waitKey(0)